#include <iostream>
#include <opencv2/opencv.hpp>
#include <string>
#include <vector>

#include "common.hpp"
#include "det/postprocess.h"

// 模拟YOLOv8输出数据
void generate_test_output(std::vector<int8_t> &output_data, int grid_h, int grid_w)
{
    int output_size = grid_h * grid_w * PROP_BOX_SIZE;
    output_data.resize(output_size, 0);

    // 模拟一些检测结果
    // 在中心位置添加一个检测框
    int center_h = grid_h / 2;
    int center_w = grid_w / 2;
    int index = (center_h * grid_w + center_w) * PROP_BOX_SIZE;

    // 设置边界框坐标 (归一化到0-1)
    output_data[index + 0] = 50;  // x center
    output_data[index + 1] = 50;  // y center
    output_data[index + 2] = 100; // width
    output_data[index + 3] = 100; // height

    // 设置类别置信度
    output_data[index + 4] = 80; // class 0 confidence (高置信度)

    // 在另一个位置添加检测框
    int pos2_h = grid_h / 4;
    int pos2_w = grid_w / 4;
    int index2 = (pos2_h * grid_w + pos2_w) * PROP_BOX_SIZE;

    output_data[index2 + 0] = 25; // x center
    output_data[index2 + 1] = 25; // y center
    output_data[index2 + 2] = 50; // width
    output_data[index2 + 3] = 50; // height
    output_data[index2 + 4] = 70; // class 0 confidence
}

int main()
{
    std::cout << "YOLOv8后处理测试程序" << std::endl;

    // 创建测试图像
    cv::Mat test_image = cv::Mat::zeros(480, 640, CV_8UC3);
    cv::rectangle(test_image, cv::Point(100, 100), cv::Point(300, 300), cv::Scalar(255, 255, 255), 2);

    // 模型参数 - 根据实际模型输入尺寸调整
    int model_in_h = 480; // 高度
    int model_in_w = 640; // 宽度
    float conf_threshold = 0.5f;
    float nms_threshold = 0.45f;

    // 生成测试输出数据
    std::vector<int8_t> output_data;
    generate_test_output(output_data, model_in_h, model_in_w);

    // 后处理参数
    BOX_RECT pads = {0, 0, 0, 0};
    float scale_w = (float)test_image.cols / model_in_w;
    float scale_h = (float)test_image.rows / model_in_h;

    // 量化参数
    std::vector<int32_t> qnt_zps = {0};
    std::vector<float> qnt_scales = {1.0f};

    // 执行后处理
    DetectionResultsGroup results;
    int ret = post_process_yolov8(output_data.data(), model_in_h, model_in_w, conf_threshold, nms_threshold, pads,
                                  scale_w, scale_h, qnt_zps, qnt_scales, &results);

    if (ret == 0)
    {
        std::cout << "后处理成功，检测到 " << results.dets.size() << " 个目标" << std::endl;

        // 在图像上绘制检测结果
        for (const auto &det : results.dets)
        {
            std::cout << "检测到目标: " << det.det_name << " 置信度: " << det.score << " 位置: (" << det.box.x << ", "
                      << det.box.y << ", " << det.box.width << ", " << det.box.height << ")" << std::endl;

            // 绘制边界框
            cv::rectangle(test_image, det.box, cv::Scalar(0, 255, 0), 2);

            // 绘制标签
            std::string label = det.det_name + " " + std::to_string(det.score).substr(0, 4);
            cv::putText(test_image, label, cv::Point(det.box.x, det.box.y - 10), cv::FONT_HERSHEY_SIMPLEX, 0.5,
                        cv::Scalar(0, 255, 0), 2);
        }

        // 显示结果
        cv::imshow("YOLOv8 Test Results", test_image);
        cv::waitKey(0);
    }
    else
    {
        std::cout << "后处理失败" << std::endl;
    }

    return 0;
}